/*-
 * Copyright (c) 2006 Allan Saddi <allan@saddi.com>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 *
 * $Id: wsgiinput.c 1506 2009-01-11 04:30:51Z Bear $
 */

#include <Python.h>

#include <assert.h>

#include "wsgi-int.h"

/* Add data as a new chunk at the tail of our chunk list */
int
wsgiBodyHandler(void *ctxt, uint8_t *data, size_t len)
{
  InputStream *self = (InputStream *)((Request *)wsgiGetRequestData(ctxt))->input;
  struct InputStream_chunk *chunk;

  if (!len) return 0;

  if ((chunk = PyMem_Malloc(sizeof(*chunk) + len)) == NULL)
    return -1;

  chunk->next = NULL;
  chunk->size = len;
  memcpy(InputStream_chunk_DATA(chunk), data, len);

  if (self->chunks) {
    self->tail->next = chunk;
    self->tail = chunk;
  }
  else {
    self->chunks = self->tail = chunk;
  }

  self->avail += len;

  return 0;
}

/* Find occurrence of a character within our buffers */
static int
InputStream_findChar(InputStream *self, int start, int c)
{
  struct InputStream_chunk *chunk = self->chunks;
  int index = 0;

  /* Find starting chunk */
  while (chunk != NULL) {
    if (chunk->size > start)
      break;

    index += chunk->size;
    start -= chunk->size;
    chunk = chunk->next;
  }

  while (chunk != NULL) {
    char *data = InputStream_chunk_DATA(chunk);
    char *found = memchr(&data[start], c, chunk->size - start);
    if (found != NULL) {
      return index + (found - data);
    }
    else {
      index += chunk->size;
      chunk = chunk->next;
      start = 0;
    }
  }

  return -1;
}

/* Consume characters between self->pos and newPos, returning it as a
   new Python string. */
static PyObject *
InputStream_consume(InputStream *self, int newPos)
{
  PyObject *result;
  int origSize = newPos - self->pos, size;
  char *data;
  struct InputStream_chunk *chunk = self->chunks, *next;
  int start = self->pos, index;

  assert(chunk != NULL);
  assert(self->pos <= chunk->size);

  if (!origSize)
    return PyString_FromString("");

  result = PyString_FromStringAndSize(NULL, origSize);
  if (result == NULL)
    return NULL;

  data = PyString_AS_STRING(result);
  index = 0;

  /* Copy chunks into string */
  while (self->pos < newPos) {
    char *cdata = InputStream_chunk_DATA(chunk);

    /* Constrain copy operations to chunk boundaries */
    size = newPos - self->pos;
    if (size > (chunk->size - start))
      size = chunk->size - start;

    memcpy(&data[index], &cdata[start], size);

    self->pos += size;
    index += size;

    /* Advance to next chunk */
    chunk = chunk->next;
    start = 0;
  }

  assert(index == origSize);
  data[index] = '\0';

  /* Free fully-consumed chunks */
  chunk = self->chunks;
  while (self->pos > chunk->size) {
    next = chunk->next;
    self->pos -= chunk->size;
    self->size -= chunk->size;
    self->avail -= chunk->size;
    PyMem_Free(chunk);
    chunk = next;
  }
  self->chunks = chunk;

  assert(self->chunks != NULL);
  assert(self->pos >= 0);
  assert(self->size > 0);
  assert(self->avail > 0);

  return result;
}

static void
InputStream_dealloc(InputStream *self)
{
  struct InputStream_chunk *chunk, *next;
  PyObject *tmp;

  chunk = self->chunks;
  while (chunk != NULL) {
    next = chunk->next;
    PyMem_Free(chunk);
    chunk = next;
  }
  self->chunks = NULL;

  tmp = (PyObject *)self->request;
  self->request = NULL;
  Py_XDECREF(tmp);

  self->ob_type->tp_free((PyObject *)self);
}

static PyObject *
InputStream_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
  InputStream *self;

  self = (InputStream *)type->tp_alloc(type, 0);
  if (self != NULL) {
    self->request = NULL;
    self->pos = 0;
    self->size = 0;
    self->avail = 0;
    self->chunks = NULL;
    self->tail = NULL;
  }

  return (PyObject *)self;
}

/* InputStream constructor. Expects to be passed the parent Request and
   the received Content-Length. */
static int
InputStream_init(InputStream *self, PyObject *args, PyObject *kwds)
{
  Request *request;
  int size;

  if (!PyArg_ParseTuple(args, "O!i", &Request_Type, &request, &size))
    return -1;

  Py_INCREF(request);
  self->request = request;
  self->size = size;

  return 0;
}

/* read() implementation */
static PyObject *
InputStream_read(InputStream *self, PyObject *args)
{
  int size = -1, newPos;

  if (!PyArg_ParseTuple(args, "|i:read", &size))
    return NULL;

  if (self->pos == self->size)
    return PyString_FromString("");

  for (;;) {
    if (size < 0 || (self->avail - self->pos) < size) {
      /* Not enough data available */
      if (self->avail == self->size) {
	/* No more coming */
	newPos = self->avail;
	break;
      }
      else {
	/* Ask for more data */
	if (wsgiGetBody(self->request->context))
	  return NULL;
	continue;
      }
    }
    else {
      newPos = self->pos + size;
      break;
    }
  }

  return InputStream_consume(self, newPos);
}

/* readline() implementation. Supports "size" argument not required by
   WSGI spec (but now required by Python 2.5's cgi module) */
static PyObject *
InputStream_readline(InputStream *self, PyObject *args)
{
  int size = -1, start, i, newPos;

  if (!PyArg_ParseTuple(args, "|i:readline", &size))
    return NULL;

  if (self->pos == self->size)
    return PyString_FromString("");

  start = self->pos;

  for (;;) {
    /* Find newline */
    i = InputStream_findChar(self, start, '\n');
    if (i < 0) {
      /* Not found? */
      if (self->avail == self->size) {
	/* No more data coming */
	newPos = self->avail;
	break;
      }
      else {
	if (size < 0 || (self->avail - self->pos) < size) {
	  /* Wait for more to come */
	  start = self->avail; /* Search new chunk only */
	  if (wsgiGetBody(self->request->context))
	    return NULL;
	  continue;
	}
	else {
	  /* Already have at least size bytes available */
	  newPos = self->pos + size;
	  break;
	}
      }
    }
    else {
      newPos = i + 1;

      /* Trim line, if necessary */
      if (size >= 0 && (self->pos + size) < newPos)
	newPos = self->pos + size;
      break;
    }
  }

  return InputStream_consume(self, newPos);
}

/* readlines() implementation. Supports "hint" argument. */
static PyObject *
InputStream_readlines(InputStream *self, PyObject *args)
{
  int hint = 0, total = 0;
  PyObject *lines = NULL, *args2 = NULL, *line;
  int len, ret;

  if (!PyArg_ParseTuple(args, "|i:readlines", &hint))
    return NULL;

  if ((lines = PyList_New(0)) == NULL)
    return NULL;

  if ((args2 = PyTuple_New(0)) == NULL)
    goto bad;

  if ((line = InputStream_readline(self, args2)) == NULL)
    goto bad;

  while ((len = PyString_GET_SIZE(line)) > 0) {
    ret = PyList_Append(lines, line);
    Py_DECREF(line);
    if (ret)
      goto bad;
    total += len;
    if (hint > 0 && total >= hint)
      break;

    if ((line = InputStream_readline(self, args2)) == NULL)
      goto bad;
  }

  Py_DECREF(line);
  Py_DECREF(args2);

  return lines;

 bad:
  Py_XDECREF(args2);
  Py_XDECREF(lines);
  return NULL;
}

/* __iter__() implementation. Simply returns self. */
static PyObject *
InputStream_iter(InputStream *self)
{
  Py_INCREF(self);
  return (PyObject *)self;
}

/* next() implementation for iteration protocol support */
static PyObject *
InputStream_iternext(InputStream *self)
{
  PyObject *line, *args;

  if ((args = PyTuple_New(0)) == NULL)
    return NULL;

  line = InputStream_readline(self, args);
  Py_DECREF(args);
  if (line == NULL)
    return NULL;

  if (PyString_GET_SIZE(line) == 0) {
    Py_DECREF(line);
    PyErr_Clear();
    return NULL;
  }

  return line;
}

static PyMethodDef InputStream_methods[] = {
  { "read", (PyCFunction)InputStream_read, METH_VARARGS,
    "Read from this input stream" },
  { "readline", (PyCFunction)InputStream_readline, METH_VARARGS,
    "Read a line from this input stream" },
  { "readlines", (PyCFunction)InputStream_readlines, METH_VARARGS,
    "Read lines from this input stream" },
  { NULL }
};

PyTypeObject InputStream_Type = {
  PyObject_HEAD_INIT(NULL)
  0,                         /*ob_size*/
  "_wsgisup.InputStream",    /*tp_name*/
  sizeof(InputStream),       /*tp_basicsize*/
  0,                         /*tp_itemsize*/
  (destructor)InputStream_dealloc, /*tp_dealloc*/
  0,                         /*tp_print*/
  0,                         /*tp_getattr*/
  0,                         /*tp_setattr*/
  0,                         /*tp_compare*/
  0,                         /*tp_repr*/
  0,                         /*tp_as_number*/
  0,                         /*tp_as_sequence*/
  0,                         /*tp_as_mapping*/
  0,                         /*tp_hash */
  0,                         /*tp_call*/
  0,                         /*tp_str*/
  0,                         /*tp_getattro*/
  0,                         /*tp_setattro*/
  0,                         /*tp_as_buffer*/
  Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_ITER, /*tp_flags*/
  "wsgi.input implementation", /* tp_doc */
  0,		             /* tp_traverse */
  0,		             /* tp_clear */
  0,		             /* tp_richcompare */
  0,		             /* tp_weaklistoffset */
  (getiterfunc)InputStream_iter, /* tp_iter */
  (iternextfunc)InputStream_iternext, /* tp_iternext */
  InputStream_methods,       /* tp_methods */
  0,                         /* tp_members */
  0,                         /* tp_getset */
  0,                         /* tp_base */
  0,                         /* tp_dict */
  0,                         /* tp_descr_get */
  0,                         /* tp_descr_set */
  0,                         /* tp_dictoffset */
  (initproc)InputStream_init, /* tp_init */
  0,                         /* tp_alloc */
  InputStream_new,           /* tp_new */
};
